# china's code:
# https://github.com/guofei9987/blind_watermark/blob/master/blind_watermark/bwm_core.py

# https://russianblogs.com/article/14572865485/
# https://pywavelets.readthedocs.io/en/latest/ref/2d-dwt-and-idwt.html#single-level-dwt2


# conda create -n DWTTEST opencv

# conda activate DWTTEST

# pip install --upgrade pip
# pip install PyWavelets
# pip install matplotlib

# conda deactivate

import cv2 as cv
import numpy as np
from pywt import dwt2, wavedec2, idwt2, waverec2
import matplotlib.pyplot as plt

im = cv.imread('img/amelieOrigin.png', 0) 
# im = cv.imread('img/amelieWM_100_000_000.png', 0)

if im.shape[1] > 500:
    (h, w) = im.shape[:2]
    width = 500
    height = int(h * (width / float(w)))
    im = cv.resize(im, (width, height), interpolation=cv.INTER_AREA)

# Single -Уровень вейвлет -разложения, доходность представляет собой низкочастотные компоненты,
#  горизонтальная высокая частота, вертикальная высокая частота и высокая частота, которые 
#  эквивалентны LL, HL, LH, HH, HH, HH на рисунке выше, соответственно.
cA, (cH, cV, cD) = dwt2(im, 'haar')
#Decomposition
cA2, (cH2, cV2, cD2), (cH1, cV1, cD1) = wavedec2(im, 'haar', level=2)

#String Каждая подраздел (низкочастотный диапазон значений CA [0,510], высокая частота [-255,255])
AH = np.concatenate([cA, cH+255], axis=1)  #axis = 1 указывает на строчки столбцов
VD = np.concatenate([cV+255, cD+255], axis=1)
res1 = np.concatenate([AH, VD], axis=0)

AH2 = np.concatenate([cA2, cH2+510], axis=1)
VD2 = np.concatenate([cV2+510, cD2+510], axis=1)
A2 = np.concatenate([AH2, VD2], axis=0)
AH1 = np.concatenate([A2, (cH1+255)*2], axis=1)
VD1 = np.concatenate([(cV1+255)*2, (cD1+255)*2], axis=1)
res2 = np.concatenate([AH1, VD1], axis=0)

#показывать
# plt.figure('2D_DWT_1level')
# plt.imshow(res1, cmap='gray', vmin=0, vmax=510)
# plt.title('1level')

# plt.figure('2D_DWT_2level')
# plt.imshow(res2, cmap='gray', vmin=0, vmax=1020)
# plt.title('2level')


# coeffs1 = cA, (cH, cV, cD)
# print("len(coeffs1[0][0]) = ", len(coeffs1[0][0])) # 250

# wm = 1_000_000_000_000_000
# for i in range(11):
# 	print(f'coeffs1[0][0][{i}] = ', coeffs1[0][0][i])
# 	coeffs1[0][0][i] = int(coeffs1[0][0][i]) + wm
# 	print(f'coeffs1[0][0][{i}] = ', coeffs1[0][0][i])
# ires1 = idwt2(coeffs1, 'haar')

# for i in range(11):
# 	print(f'coeffs1[0][{i}][{i}] = ', coeffs1[0][i][i])
# 	coeffs1[0][i][i] = int(coeffs1[0][i][i]) + wm
# 	print(f'coeffs1[0][{i}][{i}] = ', coeffs1[0][i][i])
# ires1 = idwt2(coeffs1, 'haar')

# for i in range(200, 250):
# 	print(f'coeffs1[0][{i}][{i}] = ', coeffs1[0][i][i])
# 	coeffs1[0][i][i] = int(coeffs1[0][i][i]) + wm
# 	print(f'coeffs1[0][{i}][{i}] = ', coeffs1[0][i][i])
# ires1 = idwt2(coeffs1, 'haar')


# plt.figure('iDWT_1level+WM')
# plt.imshow(ires1, cmap='gray', vmin=0, vmax=1020)
# plt.title('i1level + wm')


# coeffs2 = cA2, (cH2, cV2, cD2), (cH1, cV1, cD1)
# print('\ncoeffs2 = ', coeffs2[0][0][0])
# ires2 = waverec2(coeffs2, 'haar')

# plt.figure('iDWT_2level')
# plt.imshow(ires2, cmap='gray', vmin=0, vmax=1020)
# plt.title('i2level')

coeffs2 = cA2, (cH2, cV2, cD2), (cH1, cV1, cD1)
print('len(coeffs2[0][0] = ', len(coeffs2[0][0])) # 125
wm = 1_0
# wm = 1_000_000_000_000_000

for i in range(11):
	print(f'coeffs2[0][0][{i}] = ', coeffs2[0][0][i])
	coeffs2[0][0][i] = int(coeffs2[0][0][i]) + wm
	print(f'coeffs2[0][0][{i}] = ', coeffs2[0][0][i])

for i in range(11):
	print(f'coeffs2[0][{i}][{i}] = ', coeffs2[0][i][i])
	coeffs2[0][i][i] = int(coeffs2[0][i][i]) + wm
	print(f'coeffs2[0][{i}][{i}] = ', coeffs2[0][i][i])

for i in range(100, 125):
	print(f'coeffs2[0][{i}][{i}] = ', coeffs2[0][i][i])
	coeffs2[0][i][i] = int(coeffs2[0][i][i]) + wm
	print(f'coeffs2[0][{i}][{i}] = ', coeffs2[0][i][i])

ires2 = waverec2(coeffs2, 'haar')

plt.figure('iDWT_2level + WM')
plt.imshow(ires2, cmap='gray', vmin=0, vmax=1020)
plt.title('i2level + wm')

plt.show()